import json
import regex as re
import requests
import logging
from config import Config

logger = logging.getLogger(__name__)

def extract_json_by_key(text: str, key: str):
    """
    从文本中提取第一个包含指定键的 JSON 对象
    """
    try:
        matches = re.findall(r"\{(?:[^{}]|(?R))*\}", text)
        for match in matches:
            try:
                obj = json.loads(match)
                if key in obj:
                    logger.info(f"提取到包含键 '{key}' 的 JSON")
                    return obj
            except json.JSONDecodeError:
                continue
        logger.warning(f"未找到包含键 '{key}' 的 JSON")
        return None
    except Exception as e:
        logger.error(f"提取 JSON 失败: {e}")
        return None

def send_request(url, headers, data, stream=True, timeout=120):
    """
    通用请求函数，支持流式与非流式响应处理
    """
    if stream:
        logger.info("发起流式请求")
        try:
            response = requests.post(url, headers=headers, json=data, stream=True, timeout=timeout)
            response.raise_for_status()
        except requests.exceptions.RequestException as e:
            return {"error": f"请求失败: {str(e)}"}

        full_response = []
        with response:
            for line in response.iter_lines():
                if line:
                    try:
                        line = line.decode('utf-8')
                        if line.startswith("data: "):
                            data_str = line[6:]
                            if data_str == "[DONE]":
                                break
                            parsed = json.loads(data_str)
                            full_response.append(parsed)
                    except Exception as e:
                        logger.error(f"解析流式响应失败: {e}")
                        return {"error": f"解析流式响应失败: {str(e)}"}
        return full_response

    else:
        logger.info("发起非流式请求")
        try:
            response = requests.post(url, headers=headers, json=data, timeout=timeout)
            response.raise_for_status()
        except requests.exceptions.RequestException as e:
            return {"error": f"请求失败: {str(e)}"}

        try:
            return response.json()
        except json.JSONDecodeError:
            logger.warning("response.json() 失败，尝试手动提取 JSON")
            extracted = extract_json_by_key(response.text, key="radar_chart")
            if extracted:
                return extracted
            else:
                return {"error": "未能从响应中提取包含 'radar_chart' 的 JSON"}


def stream_main(messages,stream=True):
    url = Config.LLM_API_BASE_URL + "/chat"
    headers = {
        "Content-Type": "application/json"
    }
    data = {
        "messages": messages,
        "stream": stream
    }
    if stream:
        # 流式请求示例
        print("=== 流式请求示例 ===")
        result_stream = send_request(url, headers, data, stream=True)
        print(result_stream)

    else:

        # 非流式请求示例
        print("=== 非流式请求示例 ===")
        result_non_stream = send_request(url, headers, data, stream=False)
        print(result_non_stream)
